import collections


class Solution(object):
    def wordSquares(self, words):
        """
        :type words: List[str]
        :rtype: List[List[str]]
        """
        out = []
        all_d = []
        for i in range(4):
            d = collections.defaultdict(dict)
            for j in range(i + 1):
                d[j] = collections.defaultdict(set)
            all_d.append(d)
        for item in words:
            length = len(item)
            for i in range(length):
                all_d[length - 1][i][item[i]].add(item)

        for item in words:
            length = len(item)
            if length == 1:
                out.append([item])
            if length == 2:
                wait = all_d[1][0][item[1]]
                for w in wait:
                    out.append([item, w])
            if length == 3:
                wait0 = all_d[2][0][item[1]]
                wait1 = all_d[2][0][item[2]]
                for w in wait0:
                    wait2 = wait1 & all_d[2][1][w[2]]
                    for t in wait2:
                        out.append([item, w, t])
            if length == 4:
                wait0 = all_d[3][0][item[1]]
                wait1 = all_d[3][0][item[2]]
                wait2 = all_d[3][0][item[3]]
                for w in wait0:
                    wait1_2 = wait1 & all_d[3][1][w[2]]
                    wait2_2 = wait2 & all_d[3][1][w[3]]
                    for m in wait1_2:
                        wait2_3 = wait2_2 & all_d[3][2][m[3]]
                        for p in wait2_3:
                            out.append([item, w, m, p])
        return out


data = Solution()
words = ["area", "lead", "wall", "lady", "ball"]
print(data.wordSquares(words))
words = ["abat","baba","atan","atal"]
print(data.wordSquares(words))
